package com.rocky.component.elasticsearch;

import com.rocky.component.util.ObjectMapperUtil;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.lucene.search.function.CombineFunction;
import org.elasticsearch.common.lucene.search.function.FieldValueFactorFunction;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.BoostingQueryBuilder;
import org.elasticsearch.index.query.ConstantScoreQueryBuilder;
import org.elasticsearch.index.query.DisMaxQueryBuilder;
import org.elasticsearch.index.query.MatchQueryBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.functionscore.FieldValueFactorFunctionBuilder;
import org.elasticsearch.index.query.functionscore.FunctionScoreQueryBuilder;
import org.elasticsearch.index.query.functionscore.ScoreFunctionBuilders;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.io.IOException;

/**
 * @author rocky
 * date 2025-07-19
 * Description:
 */
@RestController
@RequestMapping("/compoundQuery")
public class CompoundQueryExample {

    private final RestHighLevelClient restHighLevelClient;

    public CompoundQueryExample(RestHighLevelClient restHighLevelClient) {
        this.restHighLevelClient = restHighLevelClient;
    }

    @GetMapping("/constantScoreQuery")
    public String constantScoreQuery() throws IOException {
        ConstantScoreQueryBuilder queryBuilder = QueryBuilders.constantScoreQuery(QueryBuilders.matchQuery("address", "Street"));

        return search(queryBuilder);
    }

    @GetMapping("/boolQuery")
    public String boolQuery() throws IOException {
        BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery()
                .must(QueryBuilders.termQuery("gender.keyword", "M"))
                .mustNot(QueryBuilders.matchQuery("address", "Lane"))
                .filter(QueryBuilders.rangeQuery("age").from(18).includeLower(true).to(36).includeUpper(true));

        return search(queryBuilder);
    }

    @GetMapping("/disMaxQuery")
    public String disMaxQuery() throws IOException {
        DisMaxQueryBuilder queryBuilder = QueryBuilders.disMaxQuery()
                .add(QueryBuilders.matchQuery("firstname", "Wolfe"))
                .add(QueryBuilders.matchQuery("address", "street"));

        return search(queryBuilder);
    }

    @GetMapping("/functionScoreQuery")
    public String functionScoreQuery() throws IOException {
        // 新得分 = 老得分 * log(1 + factor * balance)
        FieldValueFactorFunctionBuilder functionBuilder = ScoreFunctionBuilders
                .fieldValueFactorFunction("balance")
                .modifier(FieldValueFactorFunction.Modifier.LOG1P)
                .factor(0.1f);

        /**
         * 查询（query）：用于筛选出匹配的文档。
         * 函数（functions）：用于修改匹配文档的得分。
         * 得分模式（score_mode）：指定如何组合多个函数的得分。
         * 最大得分限制（max_boost）：限制最终得分的最大值。
         * 提升模式（boost_mode）：指定如何将查询得分与函数得分相结合。
         */
        FunctionScoreQueryBuilder queryBuilder = QueryBuilders.functionScoreQuery(
                        QueryBuilders.matchQuery("address", "Elton Street"),    // 基础查询
                        functionBuilder // 评分函数
                )
                .boostMode(CombineFunction.MULTIPLY)    // 评分组合方式
                .maxBoost(100f);  // 得分上限

        return search(queryBuilder);
    }

    @GetMapping("/boostingQuery")
    public String boostingQuery() throws IOException {
        MatchQueryBuilder positiveQuery = QueryBuilders.matchQuery("address", "Street");
        MatchQueryBuilder negativeQuery = QueryBuilders.matchQuery("address", "Place");
        BoostingQueryBuilder queryBuilder = QueryBuilders.boostingQuery(positiveQuery, negativeQuery)
                .negativeBoost(0.2f);

        return search(queryBuilder);
    }

    private String search(QueryBuilder queryBuilder) throws IOException {
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder().query(queryBuilder);

        SearchRequest searchRequest = new SearchRequest("bank")
                .source(searchSourceBuilder);

        SearchResponse searchResponse = restHighLevelClient.search(searchRequest, RequestOptions.DEFAULT);

        return ObjectMapperUtil.get().writeValueAsString(searchResponse.getHits());
    }
}
